from keras import Input
from keras.layers import Dense

from .base_layer import BaseLayer
from entity.model_describe import ModelDescribe
from log import log


class DenseLayer(BaseLayer):
    def check(self, model_list_line: list, models: list = None) -> bool:
        if len(model_list_line) < 2:  # 只写了Dense没有写units，报错
            log.error("传入参数不足")
            raise RuntimeError("传入参数不足")
        return False

    def transfer(self, model_list_line: list, inputs: Input, models: list = None):
        units = model_list_line[1]
        # 没有写activation，用默认relu
        activation = 'relu' if len(model_list_line) < 3 else model_list_line[2]
        return Dense(units, activation=activation)(inputs)

    def persistence(self, model_list_line: list, train_id: int) -> ModelDescribe:
        m = super().persistence(model_list_line, train_id)

        units = model_list_line[1]
        # 没有写activation，用默认relu
        activation = 'relu' if len(model_list_line) < 3 else model_list_line[2]

        m.var1 = units
        m.var2 = activation
        return m


